from contextlib import contextmanager
from threading import Lock as ThreadLock

import psycopg2

from tools.configmanager import config

"""
pg_advisory_lock (推荐):

最轻量级的解决方案
不会影响表数据
锁与事务无关（会话结束时自动释放）

咨询锁的锁ID可以是任意整数，但要在应用中统一约定
对于高并发场景，考虑设置锁超时时间
确保锁最终会被释放，避免死锁
在微服务架构中，所有服务需要使用相同的锁ID
PostgreSQL 的咨询锁性能非常高，特别适合实现分布式锁，因为它是数据库内置功能，不需要额外的维护成本。
"""


class PgDistributedLock:
    def __init__(self, dsn):
        self.dsn = dsn
        # 本地线程锁，用于单机互斥
        self.local_lock = ThreadLock()

    @contextmanager
    def global_lock(self, lock_id):
        """全局分布式锁，跨所有机器互斥"""
        conn = psycopg2.connect(self.dsn)
        try:
            cursor = conn.cursor()
            # 尝试获取全局锁
            cursor.execute("SELECT pg_try_advisory_lock(%s)", (lock_id,))
            locked = cursor.fetchone()[0]
            if not locked:
                raise Exception("Failed to acquire global distributed lock")
            try:
                yield
            finally:
                # 释放锁
                cursor.execute("SELECT pg_advisory_unlock(%s)", (lock_id,))
                conn.commit()
        finally:
            conn.close()

    @contextmanager
    def local_lock_only(self, lock_id):
        """单机互斥锁，不同机器可以并行"""
        with self.local_lock:
            conn = psycopg2.connect(self.dsn)
            try:
                cursor = conn.cursor()
                # 尝试获取本地锁（只在当前进程/线程中检查）
                cursor.execute("SELECT pg_try_advisory_lock(%s)", (lock_id,))
                locked = cursor.fetchone()[0]
                if not locked:
                    raise Exception("Failed to acquire local distributed lock")
                try:
                    yield
                finally:
                    # 释放锁
                    cursor.execute("SELECT pg_advisory_unlock(%s)", (lock_id,))
                    conn.commit()
            finally:
                conn.close()


user = config.get('node.postgresql.database.0.user')
password = config.get('node.postgresql.database.0.pwd')
host = config.get('node.postgresql.ip')
port = config.get('node.postgresql.port')
database = config.get('node.postgresql.database.0.dbname')
pg_url = f"host={host} port={port} dbname={database} user={user} password={password}"

# 创建锁管理器实例
lock_manager = PgDistributedLock(pg_url)
